#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdbool.h>

#include <time.h>
#include <unistd.h>
#include <signal.h>

#include <sys/wait.h>
#include <sys/types.h>

#include "stdbuf.h"

bool dead = false;
pid_t pid = 0;

void term(int signal)
{
    dead = true;
    if (pid) {
        kill(pid, SIGTERM);
        sleep(2);
        kill(pid, SIGKILL);
    }
    exit(0);
}

void exec(int argc, char **argv)
{
    char *args[argc + 1];
    for (int i = 0; i < argc; i++) {
        args[i] = argv[i];
    }
    args[argc] = NULL;

    char *envs[3];
    envs[0] = "LD_PRELOAD=" LIBSTDBUF_PATH;
    envs[1] = STDBUF(STDOUT_BUF, LINE_BUF);
    envs[2] = NULL;

    execvpe(argv[0], args, envs);
}

void putchar_tmsp(char c)
{
    static time_t tmsp = 0;
    if (!tmsp) {
        tmsp = time(NULL);
        char *ctime_nonewline = ctime(&tmsp);
        ctime_nonewline[strlen(ctime_nonewline) - 1] = '\0';
        printf("[%s] ", ctime_nonewline);
    }
    putchar(c);
    if (c == '\n') {
        tmsp = 0;
    }
}

int main(int argc, char **argv)
{
    if (argc == 1) {
        exit(0);
    }

    signal(SIGTERM, term);
    signal(SIGINT, term);

    while (!dead) {
        time_t start = time(NULL);

        int file_pipe[2];
        if (pipe(file_pipe) != 0) {
            fprintf(stderr, "pipe() failed!\n");
            exit(1);
        }

        pid = fork();
        if (pid == -1) {
            fprintf(stderr, "fork() failed!\n");
            exit(1);
        }
        else if (pid == 0) {
            dup2(file_pipe[1], STDOUT_FILENO);
            dup2(file_pipe[1], STDERR_FILENO);
            close(file_pipe[0]);
            close(file_pipe[1]);

            exec(argc - 1, &argv[1]);
            exit(1);
        }
        else {
            close(file_pipe[1]);

            char buffer[BUFSIZ + 1];
            while (1) {
                ssize_t size = read(file_pipe[0], buffer, BUFSIZ);
                if (size == 0 || size == -1) {
                    break;
                }
                for (int i = 0; i < size; i++) {
                    putchar_tmsp(buffer[i]);
                }
            }

            if (time(NULL) - start < 10) {
                // Suddenly, child crashed after started, retry later...
                sleep(5);
            }
        }
    }
    return 0;
}
